// Copyright (c) 2020-present,  INSPUR Co, Ltd.  All rights reserved.
// This source code is licensed under Apache 2.0 License.
#include "db/db_test_util.h"
#include "memory_arena.h"
#include "pure_mem/pmemrep.h"
#include "pure_mem/rangearena/range_arena.h"
#include "util/random.h"
#include "util/testharness.h"

namespace rocksdb {


    typedef uint64_t Key;
    void EncodeUint32(std::string* buf, uint32_t v) {
        const uint8_t tmp[sizeof(v)] = {
                uint8_t(v >> 24),
                uint8_t(v >> 16),
                uint8_t(v >> 8),
                uint8_t(v),
        };
        buf->append(reinterpret_cast<const char*>(tmp), sizeof(tmp));
    }


    void EncodeUint64(std::string* buf, uint64_t v) {
        const uint8_t tmp[sizeof(v)] = {
                uint8_t(v >> 56), uint8_t(v >> 48), uint8_t(v >> 40), uint8_t(v >> 32),
                uint8_t(v >> 24), uint8_t(v >> 16), uint8_t(v >> 8),  uint8_t(v),
        };
        buf->append(reinterpret_cast<const char*>(tmp), sizeof(tmp));
    }

    void EncodeTimestamp(std::string& s, int64_t wall_time, int32_t logical) {
        EncodeUint64(&s, uint64_t(wall_time));
        if (logical != 0) {
            EncodeUint32(&s, uint32_t(logical));
        }
    }
    const int kMVCCVersionTimestampSize = 12;

    std::string EncodeKey(const rocksdb::Slice& key, int64_t wall_time, int32_t logical) {
        std::string s;
        const bool ts = wall_time != 0 || logical != 0;
        s.reserve(key.size() + 1 + (ts ? 1 + kMVCCVersionTimestampSize : 0));
        s.append(key.data(), key.size());
        if (ts) {
            // Add a NUL prefix to the timestamp data. See DBPrefixExtractor.Transform
            // for more details.
            s.push_back(0);
            EncodeTimestamp(s, wall_time, logical);
        }
        s.push_back(char(s.size() - key.size()));
        return s;
    }

    static const char* Encode(const Key* key) {
      std::string str = std::to_string(*key);
      int len = str.length() + 8 + 1;
      uint32_t keySize = VarintLength(len) + len;
      char* ret = new char[keySize];
      memset(ret , '\0', keySize);
      char*p = EncodeVarint32(ret, len);
      memcpy(p, str.c_str(), str.length());
      return reinterpret_cast<const char*>(ret);
    }

    static Slice Decode(const char* key) {
      Slice buf = GetLengthPrefixedSlice(key);
      Slice mvccKey = Slice(buf.data(), buf.size() - 8);
      const char ts_size = mvccKey[mvccKey.size() - 1];
      if ((size_t)ts_size >= mvccKey.size()) {
        return 0;
      }
      return Slice(buf.data(), mvccKey.size() - 1 - ts_size);
    }


    typedef PureMemRep TestMemArenaRep;
    const size_t standard_key = 1000000;

    class MemArenaTest : public testing::Test {
    public:
        static void Insert(TestMemArenaRep* list, Key key) {
          char** buf = nullptr;
          void* rangeArena;
          const auto& keyy = (const Slice&) key;
          const size_t len = sizeof(Key);
          *buf = (char*)list->AllocatePure(len, buf, keyy, &rangeArena);
          memcpy(*buf, &key, sizeof(Key));
          void* buff = *buf;
          list->Insert(buff);
        }

        void Validate(TestMemArenaRep* list) {
          // Check keys exist.
          for (Key key : keys_) {
            ASSERT_TRUE(list->Contains(Encode(&key)));
          }
          // Iterate over the list, make sure keys appears in order and no extra
          // keys exist.
          InlineUserKeyIndex<const MemTableRep::KeyComparator&>::Iterator iter =
              (InlineUserKeyIndex<const MemTableRep::KeyComparator&>::Iterator)
          reinterpret_cast<const InlineUserKeyIndex<const MemTableRep::KeyComparator &> *>(list->GetIterator());
          ASSERT_FALSE(iter.Valid());
          Key zero = 0;
          iter.Seek(Encode(&zero));
          for (Key key : keys_) {
            ASSERT_TRUE(iter.Valid());
            ASSERT_EQ((char*)key, Decode(iter.key()).data());
            iter.Next();
          }
          ASSERT_FALSE(iter.Valid());
        }

    private:
        std::set<Key> keys_;
    };

    Slice genMVCCKey1(const char* key) {
      uint32_t keySize = strlen(key) + 1;
      char* ret = new char[keySize];
      memset(ret, '\0', keySize);
      memcpy(ret, key, strlen(key));
      return {ret, keySize};
    }

    //测试数据插入
    TEST_F(MemArenaTest, MemInsert) {
      const int N = 20000;
      std::set<Key> keys;
      InternalKeyComparator comparator(BytewiseComparator());
      Options options = Options();
      ImmutableCFOptions ioptions(options);
      ioptions.encode_version_node = false;
      MemArena memoryArena(comparator, ioptions);
      char tmp[1000] = {'1'};
      Slice value(tmp, 1000);
      SequenceNumber s = 100000;
      for (int i = 0; i < N; i++) {
        std::string a = std::to_string(s++) + "_test";

        Slice key = Slice(genMVCCKey1(a.data()));

        size_t key_size = key.size();
        size_t val_size = value.size();
        size_t internal_key_size = key_size + 8;

        char buf[internal_key_size];
        char* p = buf;
        memcpy(p, key.data(), key_size);
        p += key_size;
        uint64_t packed = PackSequenceAndType(s, kTypeValue);
        EncodeFixed64(p, packed);
        char ret[val_size];
        memcpy(ret, value.data(), val_size);

        Slice slice_key(buf, internal_key_size);
        Slice slice_value(ret, val_size);

        bool res = memoryArena.KVDataInsert(slice_key, slice_value);
        ASSERT_TRUE(res);

        //测试相同key导致插入失败
        if(i % 100 == 0){
          res = memoryArena.KVDataInsert(slice_key, slice_value);
          ASSERT_FALSE(res);
        }

        //相同userkey但是seq不同
        s++;
        char buf1[internal_key_size];
        p = buf1;
        memcpy(p, key.data(), key_size);
        p += key_size;
        uint64_t packed1 = PackSequenceAndType(s, kTypeValue);
        EncodeFixed64(p, packed1);
        char ret1[val_size];
        memcpy(ret1, value.data(), val_size);

        Slice slice_key1(buf1, internal_key_size);
        Slice slice_value1(ret1, val_size);

        bool res1 = memoryArena.KVDataInsert(slice_key1, slice_value1);
        ASSERT_TRUE(res1);
      }
      std::cout<<"数据insert完成."<<std::endl;

    }

    //测试数据插入，插入后拿出的key,value是否和插入前相同（测试Seek）
    TEST_F(MemArenaTest, MemInsertAndSeek) {
      const int N = 20000;
      std::set<Key> keys;
      Options options = Options();
      options.encode_version_node = false;
      ImmutableCFOptions ioptions(options);
      InternalKeyComparator comparator(BytewiseComparator());
      MemArena memoryArena(comparator, ioptions);
      char tmp[1000] = {'0'};
      Slice value(tmp, 1000);
      SequenceNumber s = 100000;
      for (int i = 0; i < N; i++) {
        std::string a = std::to_string(s++) + "_test";

        Slice key = Slice(genMVCCKey1(a.data()));

        size_t key_size = key.size();
        size_t val_size = value.size();
        size_t internal_key_size = key_size + 8;

        char buf[internal_key_size];
        char* p = buf;
        memcpy(p, key.data(), key_size);
        p += key_size;
        uint64_t packed = PackSequenceAndType(s, kTypeValue);
        EncodeFixed64(p, packed);

        char ret[val_size];
        memcpy(ret, value.data(), val_size);

        Slice slice_key(buf, internal_key_size);
        Slice slice_value(ret, val_size);


        bool insert_res = memoryArena.KVDataInsert(slice_key, slice_value);
        ASSERT_TRUE(insert_res);
        auto mem_iter = memoryArena.Seek(slice_key);     //测试内存存储区的Seek
        if(mem_iter->Valid()){
          Slice mem_key = GetLengthPrefixedSlice(mem_iter->key());
          Slice mem_value = GetLengthPrefixedSlice(mem_key.data() + mem_key.size());   //Seek单独for循环
          int res = mem_key.compare(slice_key);
          int res1 = mem_value.compare(slice_value);
          if(res != 0){
            std::cout<<"mem插入前后userkey不同， 插入mem:"<<slice_key.ToString()<<"查到："<<mem_key.ToString()<<std::endl;
            exit(1);
          }
          if(res1 != 0){
            std::cout<<"mem插入前后value不同， 插入mem:"<<slice_value.ToString()<<"查到："<<mem_value.ToString()<<std::endl;
            exit(1);
          }
        }
      }

      s = 100000;
      //Seek单独for循环
      for (int i = 0; i < N; i++) {
        std::string a = std::to_string(s++) + "_test";

        Slice key = Slice(genMVCCKey1(a.data()));

        size_t key_size = key.size();
        size_t val_size = value.size();
        size_t internal_key_size = key_size + 8;

        char buf[internal_key_size];
        char* p = buf;
        memcpy(p, key.data(), key_size);
        p += key_size;
        uint64_t packed = PackSequenceAndType(s, kTypeValue);
        EncodeFixed64(p, packed);

        char ret[val_size];
        memcpy(ret, value.data(), val_size);

        Slice slice_key(buf, internal_key_size);
        Slice slice_value(ret, val_size);


        auto mem_iter = memoryArena.Seek(slice_key);     //测试内存存储区的Seek
        if(mem_iter->Valid()){
          Slice mem_key = GetLengthPrefixedSlice(mem_iter->key());
          Slice mem_value = GetLengthPrefixedSlice(mem_key.data() + mem_key.size());
          int res = mem_key.compare(slice_key);
          int res1 = mem_value.compare(slice_value);
          if(res != 0){
            std::cout<<"mem插入前后userkey不同， 插入mem:"<<slice_key.ToString()<<"查到："<<mem_key.ToString()<<std::endl;
            exit(1);
          }
          if(res1 != 0){
            std::cout<<"mem插入前后value不同， 插入mem:"<<slice_value.ToString()<<"查到："<<mem_value.ToString()<<std::endl;
            exit(1);
          }
        }
      }

    }

    //测试数据插入(既包含kTypeValue,也有kTypeRangeDeletion)，插入后拿出的key,value是否和插入前相同（测试Seek）
    TEST_F(MemArenaTest, MemInsertAndSeek1) {
      const int N = 20000;
      std::set<Key> keys;
      Options options = Options();
      options.encode_version_node = false;
      ImmutableCFOptions ioptions(options);
      InternalKeyComparator comparator(BytewiseComparator());
      MemArena memoryArena(comparator, ioptions);
      char tmp[1000] = {'0'};
      Slice value(tmp, 1000);
      SequenceNumber s = 100000;
      ValueType type = kTypeValue;
      for (int i = 0; i < N; i++) {
        std::string a = std::to_string(s++) + "_test";

        if(s == 100100){
          std::cout<<"111"<<std::endl;
        }
        Slice key = Slice(genMVCCKey1(a.data()));

        size_t key_size = key.size();
        size_t val_size = value.size();
        size_t internal_key_size = key_size + 8;

        char buf[internal_key_size];
        char* p = buf;
        memcpy(p, key.data(), key_size);
        p += key_size;
        uint64_t packed = PackSequenceAndType(s, type);
        EncodeFixed64(p, packed);

        char ret[val_size];
        memcpy(ret, value.data(), val_size);

        Slice slice_key(buf, internal_key_size);
        Slice slice_value(ret, val_size);


        bool insert_res = memoryArena.KVDataInsert(slice_key, slice_value);
        ASSERT_TRUE(insert_res);
        auto mem_iter = memoryArena.Seek(slice_key);     //测试内存存储区的Seek
        if(mem_iter->Valid()){
          Slice mem_key = GetLengthPrefixedSlice(mem_iter->key());
          Slice mem_value = GetLengthPrefixedSlice(mem_key.data() + mem_key.size());   //Seek单独for循环
          int res = mem_key.compare(slice_key);
          int res1 = mem_value.compare(slice_value);
          if(res != 0){
            std::cout<<"mem插入前后userkey不同， 插入mem:"<<slice_key.ToString()<<"查到："<<mem_key.ToString()<<std::endl;
            exit(1);
          }
          if(res1 != 0){
            std::cout<<"mem插入前后value不同， 插入mem:"<<slice_value.ToString()<<"查到："<<mem_value.ToString()<<std::endl;
            exit(1);
          }
        }
      }

      s = 100000;
      type = kTypeRangeDeletion;
      for (int i = 0; i < 9; i++) {
        s += 100;
        std::string a = std::to_string(s-1) + "_test";

        Slice key = Slice(genMVCCKey1(a.data()));

        size_t key_size = key.size();
        size_t internal_key_size = key_size + 8;

        char buf[internal_key_size];
        char* p = buf;
        memcpy(p, key.data(), key_size);
        p += key_size;
        uint64_t packed = PackSequenceAndType(s + 5, type);
        EncodeFixed64(p, packed);

        std::string b = std::to_string(s + 5) + "_test";
        Slice key1 = Slice(b.data());
        size_t val_size = key1.size();
        char ret[val_size];
        memcpy(ret, key1.data(), key1.size());

        Slice slice_key(buf, internal_key_size);
        Slice slice_value(ret, val_size);


        bool insert_res = memoryArena.RangeDeletionInsert(slice_key, slice_value);
        ASSERT_TRUE(insert_res);
        std::cout<<"范围删除，startkey : "<<key.ToString()<<"endkey : "<<key1.ToString()<<std::endl;
        auto mem_iter = memoryArena.RangeDeletionSeek(slice_key);     //测试内存存储区的Seek
        if(mem_iter->Valid()){
          Slice mem_key = GetLengthPrefixedSlice(mem_iter->key());
          Slice mem_value = GetLengthPrefixedSlice(mem_key.data() + mem_key.size());   //Seek单独for循环
          int res = mem_key.compare(slice_key);
          int res1 = mem_value.compare(slice_value);
          std::cout<<"范围删除，mem_key : "<<mem_key.ToString()<<"mem_value : "<<mem_value.ToString()<<std::endl;
          if(res != 0){
            std::cout<<"mem插入前后userkey不同， 插入mem:"<<slice_key.ToString()<<"查到："<<mem_key.ToString()<<std::endl;
            exit(1);
          }
          if(res1 != 0){
            std::cout<<"mem插入前后value不同， 插入mem:"<<slice_value.ToString()<<"查到："<<mem_value.ToString()<<std::endl;
            exit(1);
          }
        }
      }

      s = 100000;
      //Seek单独for循环
      for (int i = 0; i < N; i++) {
        std::string a = std::to_string(s++) + "_test";

        Slice key = Slice(genMVCCKey1(a.data()));

        size_t key_size = key.size();
        size_t val_size = value.size();
        size_t internal_key_size = key_size + 8;

        char buf[internal_key_size];
        char* p = buf;
        memcpy(p, key.data(), key_size);
        p += key_size;
        uint64_t packed = PackSequenceAndType(s, kTypeValue);
        EncodeFixed64(p, packed);

        char ret[val_size];
        memcpy(ret, value.data(), val_size);

        Slice slice_key(buf, internal_key_size);
        Slice slice_value(ret, val_size);

        auto mem_iter = memoryArena.Seek(slice_key);     //测试内存存储区的Seek
        if(mem_iter == nullptr){
          std::cout<<"当前key在范围删除内, key : "<<key.ToString()<<std::endl;
        }
        if(mem_iter != nullptr && mem_iter->Valid()){
          Slice mem_key = GetLengthPrefixedSlice(mem_iter->key());
          Slice mem_value = GetLengthPrefixedSlice(mem_key.data() + mem_key.size());
          int res = mem_key.compare(slice_key);
          int res1 = mem_value.compare(slice_value);
          if(res != 0){
            std::cout<<"mem插入前后userkey不同， 插入mem:"<<slice_key.ToString()<<"查到："<<mem_key.ToString()<<std::endl;
            exit(1);
          }
          if(res1 != 0){
            std::cout<<"mem插入前后value不同， 插入mem:"<<slice_value.ToString()<<"查到："<<mem_value.ToString()<<std::endl;
            exit(1);
          }
        }
      }

    }

    //测试数据插入，统计数据量是否准确
    TEST_F(MemArenaTest, MemGetL0TotalNum) {
      const int N = 20000;
      Random rnd(10000);
      std::set<Key> keys;
      Options options = Options();
      options.encode_version_node = false;
      ImmutableCFOptions ioptions(options);
      InternalKeyComparator comparator(BytewiseComparator());
      MemArena memoryArena(comparator, ioptions);
      char tmp[1000] = {'0'};
      Slice value(tmp, 1000);
      SequenceNumber s = 100000;
      size_t total_num = 0;
      for (int i = 0; i < N; i++,s++) {
        std::string a = std::to_string(s) + "_test";

        Slice key = Slice(genMVCCKey1(a.data()));

        size_t key_size = key.size();
        size_t val_size = value.size();
        size_t internal_key_size = key_size + 8;

        char buf[internal_key_size];
        char* p = buf;
        memcpy(p, key.data(), key_size);
        p += key_size;
        uint64_t packed = PackSequenceAndType(s, kTypeValue);
        EncodeFixed64(p, packed);

        char ret[val_size];
        memcpy(ret, value.data(), val_size);

        Slice slice_key(buf, internal_key_size);
        Slice slice_value(ret, val_size);

        bool res = memoryArena.KVDataInsert(slice_key, slice_value);
        ASSERT_TRUE(res);
      }

      auto iter = memoryArena.NewMemIter(false);
      iter->SeekToFirst();
      while(iter->Valid()){
        Slice tmp_key = iter->key();
        Slice tmp_value = iter->value();
        size_t len = VarintLength(tmp_key.size()) + tmp_key.size() + VarintLength(tmp_value.size()) + tmp_value.size();
        total_num += len;
        iter->Next();
      }

      ////使用Seek查数据量
      size_t L0_num = memoryArena.GetTotalMemVolume();
      std::cout << "L0当前数据量 ：" << L0_num << std::endl;
      if (total_num != L0_num) {
        std::cout << "数据量统计出错！！ " << std::endl;
        std::cout<<"Seek到的数据量为 ："<<total_num<<std::endl;
        exit(1);
      }
    }

    //测试RangeDeletionInsert,范围删除不统计为数据量
    TEST_F(MemArenaTest, MemInsertRangeDeletion) {
      const int N = 2000;
      std::set<Key> keys;
      Options options = Options();
      options.encode_version_node = false;
      ImmutableCFOptions ioptions(options);
      InternalKeyComparator comparator(BytewiseComparator());
      MemArena memoryArena(comparator, ioptions);
      SequenceNumber s = 100000;

      for (int i = 0; i < N; i++) {
        s++;

        std::string a = std::to_string(s) + "_tcst";
        std::string b = std::to_string(s) + "_test";

        Slice key = Slice(genMVCCKey1(a.data()));
        Slice del_value = Slice(b.data(), b.size());
        size_t key_size = key.size();
        size_t del_val_size = del_value.size();
        size_t internal_key_size = key_size + 8;

        char buf[internal_key_size];
        char *p = buf;
        memcpy(p, key.data(), key_size);
        p += key_size;
        uint64_t packed = PackSequenceAndType(s, kTypeRangeDeletion);
        EncodeFixed64(p, packed);

        char ret[del_val_size];
        memcpy(ret, del_value.data(), del_val_size);

        Slice slice_key(buf, internal_key_size);
        Slice slice_value(ret, del_val_size);

        memoryArena.RangeDeletionInsert(slice_key, slice_key);
        auto mem_iter = memoryArena.RangeDeletionSeek(slice_key);  //测试内存存储区的Seek
        if (mem_iter->Valid()) {
          Slice mem_key = GetLengthPrefixedSlice(mem_iter->key());
          Slice mem_value = GetLengthPrefixedSlice(mem_key.data() + mem_key.size());
          int res = mem_key.compare(slice_key);
          int res1 = mem_value.compare(mem_value);
          if (res != 0) {
            std::cout << "mem插入前后userkey不同， 插入mem:" << slice_key.ToString() << "查到：" << mem_key.ToString()
                      << std::endl;
            exit(1);
          }

          if (res1 != 0) {
            std::cout << "mem插入前后value不同， 插入mem:" << slice_value.ToString() << "查到：" << mem_value.ToString()
                      << std::endl;
            exit(1);
          }
        }
      }
    }


    TEST_F(MemArenaTest, MemInsertAndGet) {
        InternalKeyComparator comparator(BytewiseComparator());
        Options options = Options();
        options.encode_version_node = false;
        options.rocksdb_l0_cache_flush_l1 = true;
        ImmutableCFOptions ioptions(options);
        auto* memory_arena = new MemArena(comparator, ioptions);
        char tmp[1000] = {'0'};
        Slice value(tmp, 1000);
        SequenceNumber s = 10000;
        for (int i = 0; i < 1000; i++) {
            s++;
            std::string a = std::to_string(s) + "_tast";
            std::string b = std::to_string(s) + "_tbst";
            std::string c = std::to_string(s) + "_test";
            std::string d = std::to_string(s) + "_tfst";

            std::list<std::string> strVec{a, b, c, d};
            for (const auto &str : strVec) {
                s++;
                Slice key = Slice(genMVCCKey1(str.data()));

                auto key_size = key.size();
                auto val_size = value.size();
                uint32_t internal_key_size = key_size + 8;
                const uint32_t encoded_len = VarintLength(internal_key_size) +
                                             internal_key_size + VarintLength(val_size) +
                                             val_size;

                char buf[encoded_len];
                char *p = EncodeVarint32(buf, internal_key_size);
                memcpy(p, key.data(), key_size);
                p += key_size;
                uint64_t packed = PackSequenceAndType(s, kTypeValue);
                EncodeFixed64(p, packed);
                p += 8;
                p = EncodeVarint32(p, val_size);
                memcpy(p, value.data(), val_size);

                Slice slice_key = Slice(buf + VarintLength(internal_key_size), internal_key_size);
                Slice slice_value = GetLengthPrefixedSlice(slice_key.data() + slice_key.size());
                memory_arena->KVDataInsert(slice_key, slice_value);
                std::string value_;
                Status status;
                ReadOptions read_opts;
                SequenceNumber sq = 0;
                MergeContext merge_context;
                SequenceNumber snapshot = s;
                LookupKey lkey(key, snapshot);
                bool res = memory_arena->Get(lkey, &value_, &status,
                                             &merge_context, &sq, &sq, read_opts);
                ASSERT_TRUE(res);
                ASSERT_EQ(value_.length(), 1000);
            }
        }
        std::cout << "-----------insert all records ok." << std::endl;
        std::this_thread::sleep_for(std::chrono::milliseconds(1000));



        std::cout << "delete memory_arena" << std::endl;
        delete memory_arena;
        std::cout << "-------------------------------------------------------end ..."
                  << std::endl;
    }

}  //namespace rocksdb

int main(int argc, char** argv) {
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}